
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# 一维无限深势阱： - psi''(x) = E psi,
def deriv(y,x,E):
    return [y[1], - E * y[0]]

def shoot(E):
    #print("E = ", E)
    y0 = [ 0, 1 ]
    x = np.arange(0, 1.005, 0.01) # [0,0.01,...,0.99,1]
    sol = odeint( deriv, y0, x, args=(E,) ) # E 作为参数传入 odeint
    #print("sol[-1,0]=",sol[-1,0])
    return abs(sol[-1,0])# y[0](x=1)相对目标值0的偏离量

from scipy.optimize import fsolve

root = fsolve(shoot, [0.1])
print("root=",root," pi^2 = ", np.pi*np.pi," shoot(root)=",shoot(root))

'''
作业：
V = 0.1, 0<x<0.5;
  = 0, 0.5<x<1.
psi(0)=0, psi'(0)=1, psi(1)=0. 求基态能量 E。
理论值： E = pi^2 + 0.05，其中，0.05是一阶微扰修正，不会特别准，仅供参考。
'''
